"""
# @Time    : 2021/7/1 6:53 下午
# @Author  : hezhiqiang01
# @Email   : hezhiqiang01@baidu.com
# @File    : rMAPPOPolicy.py
"""

import torch
from algorithms.algorithm.r_actor_critic import R_Actor, R_Critic
from utils.util import update_linear_schedule


class RMAPPOPolicy:
    """
    MAPPO策略类。封装演员和评论家网络以计算动作和价值函数预测。

    :param args: (argparse.Namespace) 包含相关模型和策略信息的参数。
    :param obs_space: (gym.Space) 观察空间。
    :param cent_obs_space: (gym.Space) 价值函数输入空间（MAPPO的中心化输入，IPPO的去中心化输入）。
    :param action_space: (gym.Space) 动作空间。
    :param device: (torch.device) 指定运行设备(cpu/gpu)。
    """

    def __init__(self, args, obs_space, cent_obs_space, act_space, device=torch.device("cpu")):
        self.device = device  # 运行设备
        self.lr = args.lr  # 演员学习率
        self.critic_lr = args.critic_lr  # 评论家学习率
        self.opti_eps = args.opti_eps  # 优化器epsilon参数
        self.weight_decay = args.weight_decay  # 权重衰减参数

        self.obs_space = obs_space  # 观察空间
        self.share_obs_space = cent_obs_space  # 共享观察空间
        self.act_space = act_space  # 动作空间

        # 初始化演员和评论家网络
        self.actor = R_Actor(args, self.obs_space, self.act_space, self.device)
        self.critic = R_Critic(args, self.share_obs_space, self.device)

        # 初始化演员和评论家的优化器
        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(),
                                                lr=self.lr, eps=self.opti_eps,
                                                weight_decay=self.weight_decay)
        self.critic_optimizer = torch.optim.Adam(self.critic.parameters(),
                                                 lr=self.critic_lr,
                                                 eps=self.opti_eps,
                                                 weight_decay=self.weight_decay)

    def lr_decay(self, episode, episodes):
        """
        衰减演员和评论家的学习率。
        :param episode: (int) 当前训练回合。
        :param episodes: (int) 总训练回合数。
        """
        update_linear_schedule(self.actor_optimizer, episode, episodes, self.lr)
        update_linear_schedule(self.critic_optimizer, episode, episodes, self.critic_lr)

    def get_actions(self, cent_obs, obs, rnn_states_actor, rnn_states_critic, masks, available_actions=None,
                    deterministic=False):
        """
        计算给定输入的动作和价值函数预测。
        :param cent_obs (np.ndarray): 评论家的中心化输入。
        :param obs (np.ndarray): 演员的局部智能体输入。
        :param rnn_states_actor: (np.ndarray) 如果演员是RNN，则为演员的RNN状态。
        :param rnn_states_critic: (np.ndarray) 如果评论家是RNN，则为评论家的RNN状态。
        :param masks: (np.ndarray) 表示应该重置RNN状态的点。
        :param available_actions: (np.ndarray) 表示哪些动作对智能体可用
                                  (如果为None，则所有动作都可用)
        :param deterministic: (bool) 动作是否应该是分布的模式或应该被采样。

        :return values: (torch.Tensor) 价值函数预测。
        :return actions: (torch.Tensor) 要采取的动作。
        :return action_log_probs: (torch.Tensor) 所选动作的对数概率。
        :return rnn_states_actor: (torch.Tensor) 更新后的演员网络RNN状态。
        :return rnn_states_critic: (torch.Tensor) 更新后的评论家网络RNN状态。
        """
        # 通过演员网络计算动作、动作对数概率和更新后的RNN状态
        actions, action_log_probs, rnn_states_actor = self.actor(obs,
                                                                 rnn_states_actor,
                                                                 masks,
                                                                 available_actions,
                                                                 deterministic)

        # 通过评论家网络计算价值和更新后的RNN状态
        values, rnn_states_critic = self.critic(cent_obs, rnn_states_critic, masks)
        return values, actions, action_log_probs, rnn_states_actor, rnn_states_critic

    def get_values(self, cent_obs, rnn_states_critic, masks):
        """
        获取价值函数预测。
        :param cent_obs (np.ndarray): 评论家的中心化输入。
        :param rnn_states_critic: (np.ndarray) 如果评论家是RNN，则为评论家的RNN状态。
        :param masks: (np.ndarray) 表示应该重置RNN状态的点。

        :return values: (torch.Tensor) 价值函数预测。
        """
        values, _ = self.critic(cent_obs, rnn_states_critic, masks)
        return values

    def evaluate_actions(self, cent_obs, obs, rnn_states_actor, rnn_states_critic, action, masks,
                         available_actions=None, active_masks=None):
        """
        获取用于演员更新的动作对数概率/熵和价值函数预测。
        :param cent_obs (np.ndarray): 评论家的中心化输入。
        :param obs (np.ndarray): 演员的局部智能体输入。
        :param rnn_states_actor: (np.ndarray) 如果演员是RNN，则为演员的RNN状态。
        :param rnn_states_critic: (np.ndarray) 如果评论家是RNN，则为评论家的RNN状态。
        :param action: (np.ndarray) 要计算对数概率和熵的动作。
        :param masks: (np.ndarray) 表示应该重置RNN状态的点。
        :param available_actions: (np.ndarray) 表示哪些动作对智能体可用
                                  (如果为None，则所有动作都可用)
        :param active_masks: (torch.Tensor) 表示智能体是活动的还是死亡的。

        :return values: (torch.Tensor) 价值函数预测。
        :return action_log_probs: (torch.Tensor) 输入动作的对数概率。
        :return dist_entropy: (torch.Tensor) 给定输入的动作分布熵。
        """
        # 评估动作，计算对数概率和熵
        action_log_probs, dist_entropy = self.actor.evaluate_actions(obs,
                                                                     rnn_states_actor,
                                                                     action,
                                                                     masks,
                                                                     available_actions,
                                                                     active_masks)

        # 计算价值
        values, _ = self.critic(cent_obs, rnn_states_critic, masks)
        return values, action_log_probs, dist_entropy

    def act(self, obs, rnn_states_actor, masks, available_actions=None, deterministic=False):
        """
        使用给定输入计算动作。
        :param obs (np.ndarray): 演员的局部智能体输入。
        :param rnn_states_actor: (np.ndarray) 如果演员是RNN，则为演员的RNN状态。
        :param masks: (np.ndarray) 表示应该重置RNN状态的点。
        :param available_actions: (np.ndarray) 表示哪些动作对智能体可用
                                  (如果为None，则所有动作都可用)
        :param deterministic: (bool) 动作是否应该是分布的模式或应该被采样。
        """
        # 计算动作和更新后的RNN状态
        actions, _, rnn_states_actor = self.actor(obs, rnn_states_actor, masks, available_actions, deterministic)
        return actions, rnn_states_actor
